#/*************************************************************************
#    > File Name: Param_ARCH_Operation.py
#    > Author: Yan Wang
#    > Mail: wangyan@imnu.edu.cn
#    > Created Time: 2022年10月27日 星期四 14时47分14秒
# ************************************************************************/
#!/usr/bin/env python
# -*- coding: utf-8 -*-
from Library import *
from Load_Lib import *

class Param_ARCH(object):
    ARCH_Dir={}
    ARCH_Control={}
    ARCH_Results={}
    ARCH_Setting={}

    _ARCH_path={}
    _ARCH_file={}

    _param_dir={}
    _process={}
    _MG_para={}

    event_file_data=Path("./ARCH_Param_event_file_data.dat")

    def __init__(self, param_dir, process, MG_para):
        self._param_dir=param_dir.Dir()
        self._process  =process.Var()
        self._MG_para  =MG_para.Dir()
        self.Set_Para()

# 设置基本变量
    def Set_File_Data(self, file_name):
        self.event_file_data=file_name
        self.Set_Para()

    def Set_Input(self, param_dir, process, MG_para):
        self._param_dir=param_dir.Var()
        self._process  =process.Var()
        self._MG_para  =MG_para.Dir()
        self.Set_Para()


# 设置最终输出结果
    def Set_Para(self):
        #self.Param_ARCH_Dir    (self._param_dir, self._process, self._MG_para)
        self.Param_ARCH_Dir    (self._param_dir, self._process)
        self.Param_ARCH_Control(self._param_dir, self._process, self._MG_para)


# 输出
    def Dir(self):
        return(self.ARCH_Dir)

    def Control(self):
        return(self.ARCH_Control)

    def Results(self):
        return(self.ARCH_Results)

    def Setting(self):
        return(self.ARCH_Setting)

    # 生成 ARCH 中的 Path 文件
    def Param_ARCH_Path(self, process_name,out_folder):
        param_path={"origin": {"process"     :"PROCESS          : {}_$fnum".format(process_name),
            "out_folder" :"OUTPUT_FOLDER    : "+str(out_folder),
            "tmp"        :"RECORD_FILE      : \"tmp/{}_$fnum.dat\"".format(process_name),
            "in_folder"  :"INPUT_FOLDER     : './control/control_scan_{}'".format(process_name), #这里要和上面定义的ARCH_target_control_folder对应
            "file":"INPUT_FILE              : \"file_$fnum.dat\""},
            "source" : "path.dat",
            "labels" : {"PROCESS          : process":0,
                "OUTPUT_FOLDER    : 'events/Event_Analyse_ARCH_test'":0,
                "RECORD_FILE      : 'tmp/sig_process.dat'":0, 
                "INPUT_FOLDER     : './control/control_ARCH'":0,
                "INPUT_FILE       : 'file.dat'":0}
            }
        self._ARCH_path=param_path
        return(param_path)




    def Path_Print(self):
        print("			ARCH path is : ")
        print("				origin : ")
        for key in self._ARCH_path["origin"] :
            print("					"+str(key)+" : ", self._ARCH_path["origin"][key])
        print("				source : ", self._ARCH_path["source"])
        for key in self._ARCH_path["labels"] :
            print("					"+str(key)+" : ", self._ARCH_path["labels"][key])




    # 生成 ARCH 中 File 文件中需要用得到事例的文件名及文件描述
    # 主要用于一个File文件中包含多个事例文件的情况
    def Param_ARCH_Read_MultiFile_Name_from_File(self, event_file_data):

        des=Get_File_Name_without_Ext(event_file_data)
        print("event_file_data",event_file_data,type(event_file_data))
        print("des",des)
        print("des",not event_file_Data.exists())
        if not event_file_data.exists():
            origin_file=    "FILE            : ["+"]"
            origin_file_des="FILE_DESCRIP    : ["+"]"
            origin_file_num=0
            return(origin_file, origin_file_des,origin_file_num)
        f=open(event_file_data,"r")
        lines=f.readlines()
        file_name =""
        file_des  =""
        origin_file_num=0
        for line_num, line in enumerate(lines):
            line = line.rstrip('\n')   #rstrip() 删除 string 字符串末尾的指定字符，默认为空白符，包括空格、换行符、回车符、制表符。
            if line == "":
                continue
            else:
                origin_file_num+=1

            if line_num ==0:
                file_name +=line+"\n"
                file_des  +="\""+des+"\"\n"
            else:
                file_name +="                ,"+line+"\n"
                file_des  +="                ,\""+des+"\"\n"
        origin_file=    "FILE            : ["+file_name+"]"
        origin_file_des="FILE_DESCRIP    : ["+file_des+"]"
        return(origin_file, origin_file_des,origin_file_num)

    # 生成 ARCH 中 File 文件不同情况下需要用的文件名及文件描述
    def Param_ARCH_File_Name(self, mode,event_file_data,para):
        if mode == "seperate_multievent":
            origin_file, origin_file_des, origin_file_num = self.Param_ARCH_Read_MultiFile_Name_from_File(event_file_data)

        elif mode == "seperate" or mode =="fixed_width":
            df = Load_THDM_Param(para["cut_file_name"],delimiter=",")
            file_name_list = df["#run_name"].to_list()
            file_name =     "FILE             : "
            file_name += file_name_list[0]
            file_name_list
            origin_file=    file_name
            #origin_file=    "FILE             : ["+str(para["combined_event_folder"])+"/run_$fnum/tag_1_delphes_events.root]"
            des=Get_File_Name_without_Ext(para["cut_file_name"])
            origin_file_des="FILE_DESCRIP        : [\""+des+"\"]"
            origin_file_num=1
        else:
            origin_file=""
            origin_file_des=""
            origin_file_num=1
            print("Warning: In Param_ARCH_File_Name, input mode is not correct, it should be seperate or combined: ", mode)
        return(origin_file, origin_file_des,origin_file_num)


    def Param_ARCH_File(self):
        param_file={"origin": {"filenum":"FILE_NUM        : $fNo",
            "chalnum":"CHANNEL_NUM     : $fNo",
            "ffile":"FILE            : $ffile",
            "fdes" :"FILE_DESCRIP    : $fdes",
            "cross":"CROSS_SECTION   : $fvalue"},
            "source":"file.dat",
            "labels":{"FILE_NUM        : 1":0,
                "CHANNEL_NUM     : 1":0,
                "FILE            : ['sig_Wh_4b_1_1.root']":0,
                "FILE_DESCRIP    : ['signal_1']":0,
                "CROSS_SECTION   : 32.0566":0}
            }
        self._ARCH_file=param_file
        return(param_file)



    def File_Print(self):
        print("			ARCH file is : ")
        print("				origin : ")
        for key in self._ARCH_file["origin"] :
            print("					"+str(key)+" : ", self._ARCH_file["origin"][key])
        print("				source : ", self._ARCH_file["source"])
        for key in self._ARCH_file["labels"] :
            print("					"+str(key)+" : ", self._ARCH_file["labels"][key])




    def Param_ARCH_Control(self, param_dir,process, para):
        # 设置ARCH相关参数
        self.ARCH_Control["dir"]={"source":param_dir["ARCH_template"],"target":self.ARCH_Dir["control_folder"]}

        self.Param_ARCH_Path(process["process_event_kind"],self.ARCH_Dir["event_general_folder"])
        #self.Param_ARCH_File(process["mode"], self.event_file_data, para)
        self.Param_ARCH_File()
        self.ARCH_Control["replaced_para"]=	{"path":self._ARCH_path, 
                "file": self._ARCH_file}
        return


    def Control_Print(self):
        print("	The ARCH Control is : ")
        print("		ARCH control dir : ")
        print("			source : ", self.ARCH_Control["dir"]["source"])
        print("			target : ", self.ARCH_Control["dir"]["target"])
        print("		ARCH replaced_para : ")
        self.Path_Print()
        self.File_Print()


    def Read_MultiFile_Name(self, folder_name, file_name):
        this_file=folder_name.format(file_name)
        return(this_file)

    def Replace_Para_in_ARCH(self, input_para, cut_cross):
        df_num=len(cut_cross)

        output_list={}
        df = Load_THDM_Param(self._MG_para["cut_file_name"], delimiter=",")
        cut_cross = df["cross"]
        file_name_list = df["#run_name"].to_list()
        fno=df["FileNo"]
        des=Get_File_Name_without_Ext(self._MG_para["cut_file_name"])
        for key in input_para:
            #print("key",key)
            param_list_all=[]
            for cross_num, cross_data in enumerate(cut_cross):
                format_num=str(cross_num+1).rjust(2,"0")

                ffile= file_name_list[cross_num]
                fdes=[des for i in range(fno[cross_num])]

                values={"fnum":format_num, "fNo":fno[cross_num], "fvalue":str(cross_data),"ffile":ffile, "fdes":fdes}
                #print("intput_para",input_para[key])
                input_dict=Trans_MultiValue_to_Dict(input_para[key],values)   #替换,values传参到input_para[key]中
                input_list=Trans_Dict_to_List(input_dict)   #将dict转换为list
                #print("intput_dict",input_dict)
                param_list_all.append(input_list)
            output_list[key]=param_list_all

        return(output_list)



# 生成ARCH的control文件， 
# 其中cut_cross 由Cal_Parton_Cut_Xection计算
# dir_name 是一个列表， 第一个为control文件模板，第二个为复制后使用的文件
# para_path 为修改的path的改动参数
    def Generate_Multi_Control_File(self, cut_cross):

        #从模板中复制新的control文件夹
        #print("source control", self.ARCH_Control["dir"]["source"])
        #print("target control", self.ARCH_Control["dir"]["target"])
        target_folder=Copy_File(self.ARCH_Control["dir"]["source"],self.ARCH_Control["dir"]["target"])   

# generate param files
        input_origin={}
        #print("replaced para", self.ARCH_Control["replaced_para"])
        input_replaced_para=self.ARCH_Control["replaced_para"]
        for key in input_replaced_para:
            input_origin[key] = input_replaced_para[key]["origin"]
        #print("input origin", input_origin)
        replaced_para=self.Replace_Para_in_ARCH(input_origin, cut_cross)   #替换input_origin中的cut_cross
        for key in input_replaced_para:
            File_Replace_with_Para(replaced_para[key], input_replaced_para[key]["source"], target_folder, input_replaced_para[key]["labels"])    #生成多个替换文件,替换input_para[key]["labels"]到replaced_para[key]中

        return(replaced_para)


# define general ARCH dir 
    #def Param_ARCH_Dir(self, param_dir,process, para):
    def Param_ARCH_Dir(self, param_dir,process):
        #self.ARCH_Dir["cut_file_name"]=str(para["combined_scan_file"]).format(fnum=process["BP_length"])
        #self.ARCH_Dir["nocut_file_name"]=str(para["combined_scan_file_nocut"]).format(fnum=process["BP_length"])
        self.ARCH_Dir["parallel_output_folder"]=param_dir["output_folder"] / "parallel_tmp/ARCH_results_{fname}/".format(fname=process["process_event_kind"]) / "ARCH_{}"

        self.ARCH_Dir["control_folder"]=param_dir["ARCH"] / "control/control_scan_{}".format(process["process_event_kind"])
        self.ARCH_Dir["event_general_folder"]=Path("events/Event_Analyse_scan_{fname}/".format(fname=process["process_event_kind"]))
        self.ARCH_Dir["output_event_folder"]= param_dir["ARCH"] / str(self.ARCH_Dir["event_general_folder"]).format(fname=process["process_event_kind"])
        self.ARCH_Dir["output_tmp_file"]= self.ARCH_Dir["output_event_folder"] / ("tmp/{fname}".format(fname=process["process_event_kind"]) + "_{}.dat")
        self.ARCH_Dir["output_event_file"]= self.ARCH_Dir["output_event_folder"] / ("events/{fname}".format(fname=process["process_event_kind"])+"_$fnum.root")
        # 设置BASDA相关参数，利用ARCH输出结果
        self.ARCH_Dir["summary_folder"]=param_dir["output_folder"] / "ARCH_para_{event}".format(event=process["process_event_kind"])

        self.ARCH_Dir["param_out_file"]=param_dir["output_folder"] / "param/{fname}_for_ARCH.dat".format(fname=process["process_event_kind"])

        self.ARCH_Dir["hotmap_dat_file"]=param_dir["output_folder"] / "hotmap_data/{fname}_for_ARCH_xection.dat".format(fname=process["process_event_kind"])
        self.ARCH_Dir["hotmap_C_file"]=param_dir["output_folder"] / "hotmap_code/{fname}_for_ARCH_xection.C".format(fname=process["process_event_kind"])

        self.ARCH_Setting["hotmap_variable_name"]="ARCH_xection"
        self.ARCH_Setting["plot"]={	"dat_file": self.ARCH_Dir["hotmap_dat_file"],
                "C_file": self.ARCH_Dir["hotmap_C_file"],
                "bins":process["plot"]["bins"],
                "data_length":process["BP_length"],
                "var_name": self.ARCH_Setting["hotmap_variable_name"],
                "X_name": process["plot"]["X_name"],
                "Y_name": process["plot"]["Y_name"],
                }

        for key, folder in self.ARCH_Dir.items():
            if key == "event_general_folder":
                continue
            Create_Folder(folder,"/")
        return

    def Dir_Print(self):
        print("	The ARCH Dir is : ")
        #print("		cut_file_name          : ",self.ARCH_Dir["cut_file_name"])
        #print("		nocut_file_name        : ", self.ARCH_Dir["nocut_file_name"])
        print("		control_folder         : ", self.ARCH_Dir["control_folder"])
        print("		parallel_output_folder : ", self.ARCH_Dir["parallel_output_folder"])
        print("		output_event_folder    : ", self.ARCH_Dir["output_event_folder"])
        print("		outptu_tmp_file        : ", self.ARCH_Dir["output_tmp_file"])
        print("		outptu_event_file      : ", self.ARCH_Dir["output_event_file"])
        print("		summary_folder         : ", self.ARCH_Dir["summary_folder"])
        print("		param_out_file         : ", self.ARCH_Dir["param_out_file"])
        print("		hotmap_dat_file        : ", self.ARCH_Dir["hotmap_dat_file"])
        print("		hotmap_C_file          : ", self.ARCH_Dir["hotmap_C_file"])

    def Print(self):
        print("\n ARCH Param Content is : ")
        self.Dir_Print()
        self.Control_Print()
        print("_"*50)

    def run_subprocess(self, path_file):
        ARCH_bin = self._param_dir["ARCH"] / "bin" /  "ARCH"
        result = subprocess.run([str(ARCH_bin), path_file], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        return(result)

#通过多核并行处理，执行ARCH命令，针对所有BP，计算结果
    def Run(self):
        file_general_name = "path_{}.dat"
        file_name_list=Generate_General_NameList(file_general_name,self.ARCH_Control["dir"]["target"],self._process["BP_length"])
        Create_Folder(self.ARCH_Dir["parallel_output_folder"],"/")
        Parallel_Running(self.run_subprocess, file_name_list, str(self.ARCH_Dir["parallel_output_folder"]))




# 从 ARCH 输出的tmp文件中，提取截面相关信息
# 返回值为
    def Find_ARCH_xection(self):
        ARCH_xection = []
        ARCH_number= []
        pos = []
        for file_num in range(self._process["BP_length"]):
            format_num=str(file_num+1).rjust(2,"0")
            ARCH_file = str(self.ARCH_Dir["output_tmp_file"]).format(format_num)
            #label = {'OutPut --- Efficience for Cut': 17}
            label = {'total cross section is': 11}
            word_list, input_line_list = Find_Line(ARCH_file, label)   #查找ARCH_file中的label
            strings = ''.join(word_list)
            out_file = self.ARCH_Dir["summary_folder"] /  'ARCH_output_{}.dat'.format(format_num)
            f = open(out_file, 'w')
            f.write(strings)
            f.close()
            with open(out_file, 'r') as (file):
                yfile = yaml.safe_load(file)
                strings = yfile['number passed cut basic_lep is']
                key = 'xection'
                if strings.count(key) >= 1:
                    xection = strings.split(',',3)[2].split(' ',4)[3]
                    number = strings.split(',',3)[0]
                    ARCH_xection.append(xection)
                    ARCH_number.append(number)
            pos.append(file_num)
        cut_cross=[float(value) for value in ARCH_xection]
        cut_number=[int(value) for value in ARCH_number]
        pos_number=[int(value) for value in pos]
        results= {"cross":cut_cross,"number":cut_number,"pos":pos_number}
        self.ARCH_Results=results
        return(results)


    def Collect_ARCH_Cross(self):
        origin_df=Load_THDM_Param(self._process["theoretical"],self._process["plot"]["df_para"])
        theo_df=origin_df.loc[self.ARCH_Results["pos"]]
        ARCH_df=pd.DataFrame(self.ARCH_Results)
        ARCH_df=ARCH_df.set_index("pos")
        ARCH_df=pd.concat([theo_df,ARCH_df],axis=1)
        ARCH_df.to_csv(self.ARCH_Dir["param_out_file"])
        self.ARCH_Results["df"]=ARCH_df
        return


    # 收集ARCH结果，并画出热力图
    def Get_Results(self):
        print("In ARCH:, the general ARCH results are stored in ", self.ARCH_Dir["output_tmp_file"])
        self.Find_ARCH_xection()
        self.Collect_ARCH_Cross()

        Generate_Hotmap_Data_from_Para_and_Value(self._process["theoretical"], self._process["plot"]["df_para"], self.ARCH_Results["cross"], self.ARCH_Dir["hotmap_dat_file"], self.ARCH_Setting["plot"]["var_name"])
        Generate_Hotmap_Code(self.ARCH_Setting["plot"], self._param_dir["hotmap_template"])





